#!/usr/bin/env python3
"""
交互式配置文件生成器
帮助用户创建auto_pipeline.py所需的配置文件
"""

import json
import os
from pathlib import Path
from typing import Dict, List


def get_user_input(prompt: str, default: str = "") -> str:
    """获取用户输入，支持默认值"""
    if default:
        user_input = input(f"{prompt} (默认: {default}): ").strip()
        return user_input if user_input else default
    else:
        return input(f"{prompt}: ").strip()


def validate_path(path: str) -> bool:
    """验证路径是否存在"""
    return os.path.exists(path) or '*' in path  # 支持通配符


def create_subject_config(subject_id: str) -> Dict:
    """为单个受试者创建预测文件配置"""
    print(f"\n--- 配置受试者: {subject_id} ---")
    
    config = {}
    
    # 配置proposed文件
    print("\n1. 配置proposed文件 (提出方法的预测结果)")
    add_proposed = get_user_input("是否添加proposed文件? (y/n)", "y").lower() == 'y'
    
    if add_proposed:
        config["proposed"] = []
        while True:
            print(f"\n  添加proposed文件 #{len(config['proposed']) + 1}")
            src_path = get_user_input("  源文件路径 (支持通配符)")
            if not src_path:
                break
            
            subdir = get_user_input("  子目录名称", "method1")
            
            config["proposed"].append({
                "src": src_path,
                "subdir": subdir
            })
            
            if get_user_input("  继续添加proposed文件? (y/n)", "n").lower() != 'y':
                break
    
    # 配置source文件
    print("\n2. 配置source文件 (原始图像)")
    add_source = get_user_input("是否添加source文件? (y/n)", "y").lower() == 'y'
    
    if add_source:
        config["source"] = []
        while True:
            print(f"\n  添加source文件 #{len(config['source']) + 1}")
            src_path = get_user_input("  源文件路径 (支持通配符)")
            if not src_path:
                break
            
            config["source"].append({
                "src": src_path
            })
            
            if get_user_input("  继续添加source文件? (y/n)", "n").lower() != 'y':
                break
    
    # 配置contrast_methods文件
    print("\n3. 配置contrast_methods文件 (对比方法结果)")
    add_contrast = get_user_input("是否添加contrast_methods文件? (y/n)", "y").lower() == 'y'
    
    if add_contrast:
        config["contrast_methods"] = []
        while True:
            print(f"\n  添加contrast_methods文件 #{len(config['contrast_methods']) + 1}")
            src_path = get_user_input("  源文件路径 (支持通配符)")
            if not src_path:
                break
            
            config["contrast_methods"].append({
                "src": src_path
            })
            
            if get_user_input("  继续添加contrast_methods文件? (y/n)", "n").lower() != 'y':
                break
    
    return config


def create_reference_config(subject_id: str) -> Dict:
    """为单个受试者创建参考文件配置"""
    print(f"\n--- 配置受试者参考文件: {subject_id} ---")
    
    config = {}
    
    # 配置参考图像
    print("\n1. 配置参考图像")
    add_images = get_user_input("是否添加参考图像? (y/n)", "y").lower() == 'y'
    
    if add_images:
        config["images"] = []
        while True:
            print(f"\n  添加参考图像 #{len(config['images']) + 1}")
            src_path = get_user_input("  源文件路径 (支持通配符)")
            if not src_path:
                break
            
            target_name = get_user_input("  目标文件名", "source_axial_LR_registered.nii.gz")
            
            config["images"].append({
                "src": src_path,
                "name": target_name
            })
            
            if get_user_input("  继续添加参考图像? (y/n)", "n").lower() != 'y':
                break
    
    # 配置参考掩膜
    print("\n2. 配置参考掩膜")
    add_masks = get_user_input("是否添加参考掩膜? (y/n)", "y").lower() == 'y'
    
    if add_masks:
        config["masks"] = []
        while True:
            print(f"\n  添加参考掩膜 #{len(config['masks']) + 1}")
            src_path = get_user_input("  源文件路径 (支持通配符)")
            if not src_path:
                break
            
            target_name = get_user_input("  目标文件名", "prostate_seg_source_axial_LR_registered.nii.gz")
            
            config["masks"].append({
                "src": src_path,
                "name": target_name
            })
            
            if get_user_input("  继续添加参考掩膜? (y/n)", "n").lower() != 'y':
                break
    
    return config


def main():
    """主函数"""
    print("🔧 医学图像处理流程配置文件生成器")
    print("=" * 50)
    
    # 选择配置类型
    config_type = get_user_input("选择配置类型 (1: 预测文件配置, 2: 参考文件配置, 3: 两者都创建)", "1")
    
    # 获取受试者列表
    subjects_input = get_user_input("请输入受试者ID列表 (用空格分隔)", "CHEN-REN-GENG DJY")
    subject_ids = subjects_input.split()
    
    print(f"\n将为以下受试者创建配置: {', '.join(subject_ids)}")
    
    # 创建预测文件配置
    if config_type in ["1", "3"]:
        print("\n" + "="*50)
        print("创建预测文件配置")
        print("="*50)
        
        prediction_config = {}
        for subject_id in subject_ids:
            prediction_config[subject_id] = create_subject_config(subject_id)
        
        # 保存预测配置文件
        pred_output_file = get_user_input("\n预测配置文件保存路径", "prediction_config.json")
        
        try:
            with open(pred_output_file, 'w', encoding='utf-8') as f:
                json.dump(prediction_config, f, indent=2, ensure_ascii=False)
            print(f"✅ 预测配置文件已保存: {pred_output_file}")
        except Exception as e:
            print(f"❌ 保存预测配置文件失败: {e}")
    
    # 创建参考文件配置
    if config_type in ["2", "3"]:
        print("\n" + "="*50)
        print("创建参考文件配置")
        print("="*50)
        
        reference_config = {}
        for subject_id in subject_ids:
            reference_config[subject_id] = create_reference_config(subject_id)
        
        # 保存参考配置文件
        ref_output_file = get_user_input("\n参考配置文件保存路径", "reference_config.json")
        
        try:
            with open(ref_output_file, 'w', encoding='utf-8') as f:
                json.dump(reference_config, f, indent=2, ensure_ascii=False)
            print(f"✅ 参考配置文件已保存: {ref_output_file}")
        except Exception as e:
            print(f"❌ 保存参考配置文件失败: {e}")
    
    # 提供使用示例
    print(f"\n🚀 使用示例:")
    
    if config_type in ["1", "3"]:
        pred_file = pred_output_file if config_type in ["1", "3"] else "prediction_config.json"
        print(f"# 运行完整流程（仅预测配置）")
        print(f"python auto_pipeline.py --subjects {' '.join(subject_ids)} --config {pred_file}")
    
    if config_type in ["2", "3"]:
        ref_file = ref_output_file if config_type in ["2", "3"] else "reference_config.json"
        print(f"# 运行简化流程（包含参考文件移动）")
        print(f"python simple_pipeline.py --subjects {' '.join(subject_ids)} --reference-config {ref_file}")
    
    if config_type == "3":
        print(f"# 运行完整流程（包含所有配置）")
        print(f"python auto_pipeline.py --subjects {' '.join(subject_ids)} --config {pred_output_file} --reference-config {ref_output_file}")
    
    print(f"\n# 模拟运行")
    print(f"python simple_pipeline.py --subjects {' '.join(subject_ids)} --dry-run")


if __name__ == "__main__":
    main()